import torch
from torch import Tensor
import numpy as np
import h5py
from torchsponge import Sponge
from torchsponge.potential import EnergyCell
from torchsponge.potential.pyscfcell import PySCFForceCell
from torchsponge import Protein, WithForceCell, UpdaterMD, ForceField, WithEnergyCell, RunOneStepCell, ForceFieldBase
from torchsponge.callback import WriteH5MD
from torchsponge.sampling import Metadynamics
from torchsponge.function import VelocityGenerator
from torchsponge.colvar import Torsion
from torchsponge.function import PI

mol = Protein(pdb='alad.pdb')
potential = ForceField(mol, 'AMBER.FF14SB')
phi = Torsion([4, 6, 8, 14])
psi = Torsion([6, 8, 14, 16])

min_opt = torch.optim.Adam(mol.parameters(), 1e-4)
mini = Sponge(mol, potential, min_opt, metrics={'phi': phi, 'psi': psi})
mini.run(100, 10)

class BlankEnergy(EnergyCell):
    def __init__(self, name: str = 'energy', length_unit: str = 'nm', energy_unit: str = 'kj/mol', use_pbc: bool = None, **kwargs):
        super().__init__(name=name, length_unit=length_unit, energy_unit=energy_unit, use_pbc=use_pbc, **kwargs)

    def forward(self, coordinate: Tensor, neighbour_index: Tensor = None, neighbour_mask: Tensor = None, 
                neighbour_vector: Tensor = None, neighbour_distance: Tensor = None, pbc_box: Tensor = None):
        b = coordinate.shape[0]
        energy = torch.zeros(b, 1, dtype=torch.float32, device=coordinate.device)
        return energy

metad = Metadynamics(
    colvar=[phi, psi],
    update_pace=10,
    height=2.5,
    sigma=0.05,
    grid_min=-PI,
    grid_max=PI,
    grid_bin=360,
    temperature=300,
    bias_factor=100,
)

blank = BlankEnergy()
ff = ForceFieldBase(blank)
bs = WithEnergyCell(mol,ff,bias=metad)
pyscf_frc = PySCFForceCell(system=mol, pyscf_option='ks', mol_parameters={'basis': '6-31g','verbose':0}, ks_parameters={'xc':'b3lyp'}, use_pbc=False)
wf = WithForceCell(mol, pyscf_frc)

vgen = VelocityGenerator(300)
velocity = vgen(mol.shape, mol.atom_mass)
opt = UpdaterMD(
    mol,
    temperature = 300,
    integrator='leap_frog',
    thermostat='langevin',
    time_step=2e-3,
    velocity=velocity
)
onestep = RunOneStepCell(energy=bs, force=wf, optimizer=opt)
md = Sponge(network=onestep, write_h5md=WriteH5MD(mol, 'test.h5md', 100, write_metrics=True, write_bias=True), metrics={'phi': phi, 'psi': psi})
md.run(100,10,10)  # total steps, print_interval, write_info_interval

with h5py.File('test.h5md') as f:
    md_hills = np.array(f['parameters']['meta_potential']['hills0'])
    md_grids = np.array(f['parameters']['meta_potential']['grids0'])

print(md_hills)
print(md_grids)